//! 快照管理
//! 
//! Firecracker VM快照的创建、存储和恢复

use crate::{Result, Error};
use crate::types::{SandboxId, Language};
use std::path::PathBuf;
use std::time::SystemTime;
use serde::{Serialize, Deserialize};
use tokio::fs;
use std::collections::HashMap;

/// 快照元数据
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SnapshotMetadata {
    /// 快照ID
    pub id: String,
    /// 关联的沙箱ID
    pub sandbox_id: SandboxId,
    /// 语言类型
    pub language: Language,
    /// 创建时间
    pub created_at: SystemTime,
    /// 快照大小（字节）
    pub size_bytes: u64,
    /// 快照文件路径
    pub snapshot_path: PathBuf,
    /// 内存文件路径
    pub memory_path: PathBuf,
    /// 快照路径（兼容性字段）
    pub path: String,
    /// 快照类型
    pub snapshot_type: SnapshotType,
    /// 是否已压缩
    pub compressed: bool,
}

/// 快照类型
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SnapshotType {
    /// 完整快照
    Full,
    /// 增量快照
    Incremental,
    /// 语言环境快照
    Language,
}

/// 快照管理器
pub struct SnapshotManager {
    /// 快照存储根目录
    storage_root: PathBuf,
    /// 快照元数据缓存
    metadata_cache: HashMap<String, SnapshotMetadata>,
}

impl SnapshotManager {
    /// 创建快照管理器
    pub fn new(storage_root: PathBuf) -> Self {
        Self { 
            storage_root,
            metadata_cache: HashMap::new(),
        }
    }
    
    /// 创建快照
    pub async fn create_snapshot(
        &mut self,
        sandbox_id: SandboxId,
        language: Language,
        snapshot_type: SnapshotType,
    ) -> Result<SnapshotMetadata> {
        let snapshot_id = format!("{}_{}_{}", language, sandbox_id, uuid::Uuid::new_v4());
        let snapshot_path = self.storage_root.join(format!("{}.snapshot", snapshot_id));
        let memory_path = self.storage_root.join(format!("{}.memory", snapshot_id));
        
        tracing::info!("Creating snapshot: {} for sandbox: {}", snapshot_id, sandbox_id);
        
        // 确保存储目录存在
        fs::create_dir_all(&self.storage_root).await
            .map_err(|e| Error::firecracker(format!("Failed to create storage directory: {}", e)))?;
        
        // 创建快照文件
        self.create_snapshot_files(&snapshot_path, &memory_path, &snapshot_type).await?;
        
        // 计算文件大小
        let size_bytes = self.calculate_snapshot_size(&snapshot_path, &memory_path).await?;
        
        let metadata = SnapshotMetadata {
            id: snapshot_id.clone(),
            sandbox_id,
            language,
            created_at: SystemTime::now(),
            size_bytes,
            snapshot_path: snapshot_path.clone(),
            memory_path: memory_path.clone(),
            path: snapshot_path.to_string_lossy().to_string(),
            snapshot_type,
            compressed: false,
        };
        
        // 持久化快照元数据
        self.save_metadata(&metadata).await?;
        
        // 缓存元数据
        self.metadata_cache.insert(snapshot_id, metadata.clone());
        
        tracing::info!("Snapshot created successfully: {}", metadata.id);
        Ok(metadata)
    }
    
    /// 从快照恢复VM
    pub async fn restore_from_snapshot(&self, snapshot_id: &str) -> Result<SandboxId> {
        tracing::info!("Restoring from snapshot: {}", snapshot_id);
        
        // 加载快照元数据
        let metadata = self.load_metadata(snapshot_id).await?;
        
        // 验证快照文件完整性
        self.verify_snapshot_files(&metadata).await?;
        
        // 准备恢复环境
        self.prepare_restore_environment(&metadata).await?;
        
        tracing::info!("Snapshot restored successfully: {}", snapshot_id);
        Ok(metadata.sandbox_id)
    }
    
    /// 列出所有快照
    pub async fn list_snapshots(&mut self) -> Result<Vec<SnapshotMetadata>> {
        tracing::debug!("Listing all snapshots");
        
        // 扫描存储目录中的所有.meta文件
        let mut snapshots = Vec::new();
        let mut entries = fs::read_dir(&self.storage_root).await
            .map_err(|e| Error::firecracker(format!("Failed to read storage directory: {}", e)))?;
        
        while let Some(entry) = entries.next_entry().await
            .map_err(|e| Error::firecracker(format!("Failed to read directory entry: {}", e)))? {
            
            let path = entry.path();
            if let Some(extension) = path.extension() {
                if extension == "meta" {
                    if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
                        match self.load_metadata(stem).await {
                            Ok(metadata) => {
                                self.metadata_cache.insert(stem.to_string(), metadata.clone());
                                snapshots.push(metadata);
                            }
                            Err(e) => {
                                tracing::warn!("Failed to load snapshot metadata {}: {}", stem, e);
                            }
                        }
                    }
                }
            }
        }
        
        // 按创建时间排序
        snapshots.sort_by(|a, b| b.created_at.cmp(&a.created_at));
        
        tracing::info!("Found {} snapshots", snapshots.len());
        Ok(snapshots)
    }
    
    /// 删除快照
    pub async fn delete_snapshot(&mut self, snapshot_id: &str) -> Result<()> {
        tracing::info!("Deleting snapshot: {}", snapshot_id);
        
        let metadata = self.load_metadata(snapshot_id).await?;
        
        // 删除快照文件
        if metadata.snapshot_path.exists() {
            fs::remove_file(&metadata.snapshot_path).await
                .map_err(|e| Error::firecracker(format!("Failed to delete snapshot file: {}", e)))?;
        }
        
        // 删除内存文件
        if metadata.memory_path.exists() {
            fs::remove_file(&metadata.memory_path).await
                .map_err(|e| Error::firecracker(format!("Failed to delete memory file: {}", e)))?;
        }
        
        // 删除元数据文件
        let metadata_path = self.storage_root.join(format!("{}.meta", snapshot_id));
        if metadata_path.exists() {
            fs::remove_file(&metadata_path).await
                .map_err(|e| Error::firecracker(format!("Failed to delete metadata file: {}", e)))?;
        }
        
        // 从缓存中移除
        self.metadata_cache.remove(snapshot_id);
        
        tracing::info!("Snapshot deleted successfully: {}", snapshot_id);
        Ok(())
    }
    
    /// 压缩快照（节省存储空间）
    pub async fn compress_snapshot(&mut self, snapshot_id: &str) -> Result<()> {
        tracing::info!("Compressing snapshot: {}", snapshot_id);
        
        let mut metadata = self.load_metadata(snapshot_id).await?;
        
        if metadata.compressed {
            tracing::info!("Snapshot {} is already compressed", snapshot_id);
            return Ok(());
        }
        
        // 压缩快照文件
        let compressed_snapshot = self.compress_file(&metadata.snapshot_path).await?;
        let compressed_memory = self.compress_file(&metadata.memory_path).await?;
        
        // 替换原文件
        fs::rename(&compressed_snapshot, &metadata.snapshot_path).await
            .map_err(|e| Error::firecracker(format!("Failed to replace snapshot file: {}", e)))?;
        
        fs::rename(&compressed_memory, &metadata.memory_path).await
            .map_err(|e| Error::firecracker(format!("Failed to replace memory file: {}", e)))?;
        
        // 更新元数据
        metadata.compressed = true;
        metadata.size_bytes = self.calculate_snapshot_size(&metadata.snapshot_path, &metadata.memory_path).await?;
        
        self.save_metadata(&metadata).await?;
        self.metadata_cache.insert(snapshot_id.to_string(), metadata);
        
        tracing::info!("Snapshot compressed successfully: {}", snapshot_id);
        Ok(())
    }
    
    /// 验证快照完整性
    pub async fn verify_snapshot(&self, snapshot_id: &str) -> Result<bool> {
        tracing::info!("Verifying snapshot integrity: {}", snapshot_id);
        
        let metadata = self.load_metadata(snapshot_id).await?;
        
        // 验证快照文件存在且可读
        if !metadata.snapshot_path.exists() {
            tracing::error!("Snapshot file missing: {:?}", metadata.snapshot_path);
            return Ok(false);
        }
        
        if !metadata.memory_path.exists() {
            tracing::error!("Memory file missing: {:?}", metadata.memory_path);
            return Ok(false);
        }
        
        // 验证文件大小一致性
        let current_size = self.calculate_snapshot_size(&metadata.snapshot_path, &metadata.memory_path).await?;
        if current_size != metadata.size_bytes {
            tracing::error!("Snapshot size mismatch: expected {}, got {}", metadata.size_bytes, current_size);
            return Ok(false);
        }
        
        // 验证文件完整性（简单的可读性检查）
        match self.verify_snapshot_files(&metadata).await {
            Ok(_) => {
                tracing::info!("Snapshot verification successful: {}", snapshot_id);
                Ok(true)
            }
            Err(e) => {
                tracing::error!("Snapshot verification failed: {}", e);
                Ok(false)
            }
        }
    }
    
    /// 创建快照文件
    async fn create_snapshot_files(
        &self,
        snapshot_path: &PathBuf,
        memory_path: &PathBuf,
        snapshot_type: &SnapshotType,
    ) -> Result<()> {
        // 模拟快照创建过程
        // 在真实实现中，这里会调用Firecracker API创建快照
        
        // 创建模拟快照文件
        let snapshot_content = format!("FIRECRACKER_SNAPSHOT_{:?}", snapshot_type);
        fs::write(snapshot_path, snapshot_content.as_bytes()).await
            .map_err(|e| Error::firecracker(format!("Failed to write snapshot file: {}", e)))?;
        
        // 创建模拟内存文件
        let memory_content = vec![0u8; 1024]; // 1KB模拟内存数据
        fs::write(memory_path, memory_content).await
            .map_err(|e| Error::firecracker(format!("Failed to write memory file: {}", e)))?;
        
        // 模拟创建延迟
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
        
        Ok(())
    }
    
    /// 计算快照总大小
    async fn calculate_snapshot_size(&self, snapshot_path: &PathBuf, memory_path: &PathBuf) -> Result<u64> {
        let snapshot_size = fs::metadata(snapshot_path).await
            .map_err(|e| Error::firecracker(format!("Failed to read snapshot metadata: {}", e)))?
            .len();
        
        let memory_size = fs::metadata(memory_path).await
            .map_err(|e| Error::firecracker(format!("Failed to read memory metadata: {}", e)))?
            .len();
        
        Ok(snapshot_size + memory_size)
    }
    
    /// 验证快照文件
    async fn verify_snapshot_files(&self, metadata: &SnapshotMetadata) -> Result<()> {
        // 尝试读取文件开头确保可读性
        let mut buffer = [0u8; 32];
        
        let mut snapshot_file = fs::File::open(&metadata.snapshot_path).await
            .map_err(|e| Error::firecracker(format!("Cannot open snapshot file: {}", e)))?;
        
        let mut memory_file = fs::File::open(&metadata.memory_path).await
            .map_err(|e| Error::firecracker(format!("Cannot open memory file: {}", e)))?;
        
        // 读取文件头部验证格式
        use tokio::io::AsyncReadExt;
        snapshot_file.read(&mut buffer).await
            .map_err(|e| Error::firecracker(format!("Cannot read snapshot file: {}", e)))?;
        
        memory_file.read(&mut buffer).await
            .map_err(|e| Error::firecracker(format!("Cannot read memory file: {}", e)))?;
        
        Ok(())
    }
    
    /// 准备恢复环境
    async fn prepare_restore_environment(&self, metadata: &SnapshotMetadata) -> Result<()> {
        // 验证恢复所需的资源和权限
        // 在真实实现中，这里会准备VM恢复所需的环境
        
        tracing::debug!("Preparing restore environment for snapshot: {}", metadata.id);
        
        // 模拟环境准备
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
        
        Ok(())
    }
    
    /// 压缩文件
    async fn compress_file(&self, file_path: &PathBuf) -> Result<PathBuf> {
        let compressed_path = file_path.with_extension("gz");
        
        // 在真实实现中，这里会使用实际的压缩算法（如gzip）
        // 为了简化，我们只是复制文件并添加压缩标识
        
        let content = fs::read(file_path).await
            .map_err(|e| Error::firecracker(format!("Failed to read file for compression: {}", e)))?;
        
        // 模拟压缩：在内容前添加压缩头
        let mut compressed_content = b"GZIP_HEADER".to_vec();
        compressed_content.extend_from_slice(&content);
        
        fs::write(&compressed_path, compressed_content).await
            .map_err(|e| Error::firecracker(format!("Failed to write compressed file: {}", e)))?;
        
        Ok(compressed_path)
    }
    
    /// 保存快照元数据
    async fn save_metadata(&self, metadata: &SnapshotMetadata) -> Result<()> {
        let metadata_path = self.storage_root.join(format!("{}.meta", metadata.id));
        let json = serde_json::to_string_pretty(metadata)
            .map_err(|e| Error::firecracker(format!("Failed to serialize metadata: {}", e)))?;
        
        fs::write(metadata_path, json).await
            .map_err(|e| Error::firecracker(format!("Failed to write metadata: {}", e)))?;
        
        Ok(())
    }
    
    /// 加载快照元数据
    async fn load_metadata(&self, snapshot_id: &str) -> Result<SnapshotMetadata> {
        // 先检查缓存
        if let Some(metadata) = self.metadata_cache.get(snapshot_id) {
            return Ok(metadata.clone());
        }
        
        let metadata_path = self.storage_root.join(format!("{}.meta", snapshot_id));
        let json = fs::read_to_string(metadata_path).await
            .map_err(|e| Error::firecracker(format!("Failed to read metadata: {}", e)))?;
        
        let metadata: SnapshotMetadata = serde_json::from_str(&json)
            .map_err(|e| Error::firecracker(format!("Failed to parse metadata: {}", e)))?;
        
        Ok(metadata)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    use uuid::Uuid;
    
    #[tokio::test]
    async fn test_snapshot_creation() {
        let temp_dir = tempdir().unwrap();
        let mut manager = SnapshotManager::new(temp_dir.path().to_path_buf());
        
        let sandbox_id = Uuid::new_v4();
        let language = Language::Python;
        
        let metadata = manager.create_snapshot(sandbox_id, language, SnapshotType::Full).await.unwrap();
        
        assert_eq!(metadata.sandbox_id, sandbox_id);
        assert_eq!(metadata.language, Language::Python);
        assert!(metadata.size_bytes > 0);
    }
    
    #[tokio::test]
    async fn test_snapshot_restoration() {
        let temp_dir = tempdir().unwrap();
        let mut manager = SnapshotManager::new(temp_dir.path().to_path_buf());
        
        let sandbox_id = Uuid::new_v4();
        let metadata = manager.create_snapshot(sandbox_id, Language::Python, SnapshotType::Full).await.unwrap();
        
        let restored_id = manager.restore_from_snapshot(&metadata.id).await.unwrap();
        assert_eq!(restored_id, sandbox_id);
    }
}